package com.ruijie.camerax.media;

import android.util.Log;

import androidx.camera.core.ImageProxy;

import java.nio.ByteBuffer;

/**
 * yuv420p:  yyyyyyyyuuvv
 * yuv420sp: yyyyyyyyuvuv
 * nv21:     yyyyyyyyvuvu
 */

public class Image2yuvUtils {
    public static final int YUV420P = 0;
    public static final int YUV420SP = 1;
    public static final int NV21 = 2;
    private static final String TAG = "Image2yuvUtils";
    private static final int LOG_MAX_COUNT = 5;
    private static int log_count = 0;

    /**
     * //cameraX生成的是YUV_420_888
     *
     * @param image
     * @param yuv420Frame
     * @return
     */
    public static boolean image2I420p(ImageProxy image, YUV420Frame yuv420Frame) {
        final int width = image.getWidth();
        final int height = image.getHeight();


        int yPixelCount = width * height;
        int uvPixelCount = yPixelCount / 4;
        int totalCount = 3 * yPixelCount / 2;

        final int uvWidth = width / 2;
        final int uvHeight = height / 2;

        if (yuv420Frame.yuvBytes.length < totalCount) {
            throw new RuntimeException("yuvBytes length must > w*h*3/2");
        }
        //帧时间戳
        yuv420Frame.timestamp = image.getImageInfo().getTimestamp();

        ImageProxy.PlaneProxy[] planes = image.getPlanes();

        if (log_count++ < LOG_MAX_COUNT) {
            Log.d(TAG, "image2I420p : " + image.getFormat() + " / " + width + ":" + height + " ; pixelStride1 = " + planes[1].getPixelStride() + ";rowStride = " + planes[1].getRowStride());
        }

        //填充Y分量
        fillPlaneDataRow2Row(yuv420Frame.yuvBytes, 0, planes[0], width, height);

        int pixelStride2 = planes[1].getPixelStride();
        //y数据的这个值只能是：1
        // pixelStride2 = 1. planes[1] = {UUUU...}，planes[2] = {VVVV...};  //I420
        // pixelStride2 = 2. planes[1] = {UVUV...}，planes[2] = {VUVU...}。
        //填充UV分量   此处是按照YUV420Sp格式存储
        int uStartIndex = yPixelCount;
        int vStartIndex = yPixelCount + uvPixelCount;

        //=============
        if (pixelStride2 == 1) {
            // u分量
            fillPlaneDataRow2Row(yuv420Frame.yuvBytes, uStartIndex, planes[1], uvWidth, uvHeight);
            //v分量
            fillPlaneDataRow2Row(yuv420Frame.yuvBytes, vStartIndex, planes[2], uvWidth, uvHeight);
        } else {
            // uv分量
            ImageProxy.PlaneProxy uvImage = planes[1];
            ByteBuffer uvBuffer = uvImage.getBuffer();

            // vu分量
            ImageProxy.PlaneProxy vuImage = planes[2];
            ByteBuffer vuBuffer = vuImage.getBuffer();

            Log.d(TAG, "uvBuffer: " + uvBuffer.remaining());
            for (int i = 0; i < uvHeight; i++) {
                for (int j = 0; j < uvWidth; j++) {
                    //u值
                    yuv420Frame.yuvBytes[uStartIndex + (i * uvWidth) + j] = uvBuffer.get();
                    if (uvBuffer.hasRemaining()) uvBuffer.get();
                    //v值
                    yuv420Frame.yuvBytes[vStartIndex + (i * uvWidth) + j] = vuBuffer.get();
                    if (vuBuffer.hasRemaining()) vuBuffer.get();
                }
            }
        }
        return true;
    }


    /**
     * //cameraX生成的是YUV_420_888
     *
     * @param image
     * @param yuv420Frame
     * @return
     */
    public static boolean image2NV21(ImageProxy image, YUV420Frame yuv420Frame) {
        final int width = image.getWidth();
        final int height = image.getHeight();

        int yPixelCount = width * height;
        int uvPixelCount = yPixelCount / 4;
        int totalCount = 3 * yPixelCount / 2;

        final int uvWidth = width / 2;
        final int uvHeight = height / 2;

        if (yuv420Frame.yuvBytes.length < totalCount) {
            throw new RuntimeException("yuvBytes length must > w*h*3/2");
        }
        //帧时间戳
        yuv420Frame.timestamp = image.getImageInfo().getTimestamp();

        ImageProxy.PlaneProxy[] planes = image.getPlanes();

        if (log_count++ < LOG_MAX_COUNT) {
            String builder = "image2NV21 : " +
                    "format = " + image.getFormat() + ";" +
                    "width = " + width + ";" +
                    "height = " + height + ";" +
                    "plane[0].size = " + planes[0].getBuffer().remaining() + ";" +
                    "plane[1].size = " + planes[1].getBuffer().remaining() + ";" +
                    "pixelStride1 = " + planes[1].getPixelStride() + ";" +
                    "rowStride1 = " + planes[1].getRowStride() + ";";
            Log.d(TAG, builder);
        }


        //填充Y分量
        fillPlaneDataRow2Row(yuv420Frame.yuvBytes, 0, planes[0], width, height);


        int pixelStride1 = planes[1].getPixelStride();
        int uStartIndex = yPixelCount;
        int vStartIndex = yPixelCount + uvPixelCount;

        //=============UV===
        if (pixelStride1 == 1) {
            // u分量
            ImageProxy.PlaneProxy uImage = planes[1];
            ByteBuffer uBuffer = uImage.getBuffer();

            //v分量
            ImageProxy.PlaneProxy vImage = planes[2];
            ByteBuffer vBuffer = vImage.getBuffer();

            if (uImage.getRowStride() == uvWidth) {
                for (int i = 0; i < uvHeight; i++) {
                    for (int j = 0; j < width; j += 2) {
                        yuv420Frame.yuvBytes[uStartIndex + i * width + j] = uBuffer.get();
                        yuv420Frame.yuvBytes[uStartIndex + i * width + j + 1] = vBuffer.get();
                    }
                }
            } else {
                byte[] skipBytes = new byte[uImage.getRowStride() - uvWidth];
                for (int i = 0; i < uvHeight; i++) {
                    for (int j = 0; j < width; j += 2) {
                        yuv420Frame.yuvBytes[uStartIndex + i * width + j] = uBuffer.get();
                        yuv420Frame.yuvBytes[uStartIndex + i * width + j + 1] = vBuffer.get();
                    }
                    if (i < uvHeight - 1) {
                        uBuffer.get(skipBytes);
                        vBuffer.get(skipBytes);
                    }
                }
            }

        } else {
            // uv分量  这里的planes[1]表示的uv分量最后少一个字节的v分量，为了快速取出uv分量，这里舍弃一个v不过并不影响
            fillPlaneDataRow2Row(yuv420Frame.yuvBytes, uStartIndex, planes[1], width, uvHeight);
        }
        return true;
    }

    private static void fillPlaneDataRow2Row(byte[] destBytes, int offset, ImageProxy.PlaneProxy plane, int width, int height) {
        // uv分量  这里的planes[1]表示的uv分量最后少一个字节的v分量，为了快速取出uv分量，这里舍弃一个v不过并不影响
        ByteBuffer buffer = plane.getBuffer();
        if (plane.getRowStride() == width) {
            buffer.get(destBytes, offset, buffer.remaining());
        } else {
            byte[] skipBytes = new byte[plane.getRowStride() - width];
            for (int i = 0; i < height; i++) {
                if (i == height - 1) {
                    //最后一行没有填充,并且数据长度有可能<width
                    buffer.get(destBytes, offset + i * width, buffer.remaining());
                } else {
                    buffer.get(destBytes, offset + i * width, width);
                    //跳过填充的0000数据
                    buffer.get(skipBytes);
                }
            }
        }
    }

    private static void fillPlaneDataByte2Byte() {

    }
}
